
import requests
from conn import configs
import datetime
headers = {
    'Authorization': f'Bearer {configs.kb_api_key}',
    'Content-Type': 'application/json'
}
news_dataset_id = 'd0ae1e81-82f0-4878-9ba6-be213842983c'
metas = {'title_zh': '02cefeb4-2979-46e4-9c51-0b9a3093021f',
         'url': 'd22d07c9-dc31-4edd-b669-3e685f3a435c',
         'created': '4dcda3b5-a5e8-4fb6-b43e-33f583e0a384'}


def add(text,name):

    data = {
        "name": name,
        "text": f"新闻名:{name},正文:{text}",
        "indexing_technique": "high_quality",
        "process_rule": {
            "mode": "automatic"
        }
    }
    url = f'http://{configs.dify_host}/v1/datasets/{news_dataset_id}/document/create-by-text'
    res = requests.post(url=url,headers=headers,json=data)
    if res.status_code==200:
        return res.json()['document']['id']

def add_metadata(doc_id,metadata):
    print( metadata)
    url = f'http://{configs.dify_host}/v1/datasets/{news_dataset_id}/documents/metadata'
    metadata_list = [
        {
            "id": metas[key],
            "value": metadata[key],
            "name": key
        }
        for key in metas
    ]
    data = {
        "operation_data": [
            {
                "document_id": doc_id,
                "metadata_list": metadata_list
            }
        ]
    }
    r = requests.post(url=url,headers=headers,json=data)
    print(r.text)

def add_news(news):
    content = news['title_zh'] if news['is_protected'] else news['content_zh']
    doc_id = add(content,news['title_zh'])
    # 字符串时间变成时间
    news['created'] = datetime.datetime.strptime(news['created'], "%Y-%m-%dT%H:%M:%S").strftime('%Y-%m-%d')
    add_metadata(doc_id,news)

def add_news_list(news_list):
    for news in news_list:
        add_news(news)

def list_metadata():
    r = requests.get(url=f'http://{configs.dify_host}/v1/datasets/{news_dataset_id}/metadata',headers=headers)
    datas = r.json()['doc_metadata']
    meta = {d['name']:d['id'] for d in datas}
    print(meta)

def search(query):
    url = f'http://{configs.dify_host}/v1/datasets/{news_dataset_id}/retrieve'
    datas = {
        "query": query,
        "top_k": 10
    }
    r = requests.post(url=url,headers=headers,json=datas)
    records =  r.json()['records']
    return [record['segment']['document']['doc_metadata']['url'] for record in records]


if __name__ == '__main__':
    #list_metadata()
    a = search('AI')
    print(a)

